import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import os
from pathlib import Path
from src.utils.paths import RAW_DATA_DIR, PROCESSED_DATA_DIR
from sklearn.preprocessing import LabelEncoder
from sklearn.base import BaseEstimator, TransformerMixin

def feature_engineering(data):
    '''
    对给定的数据源，进行特征工程处理，提取出关键的特征
    :param data:数据源
    :return:完整数据，特征列名
    '''
    feature_data = data.copy()
    # 是否管理岗位 = 工作角色 ∈ ['Manager', 'Director', 'Executive']
    feature_data['isManager'] = feature_data['JobRole_Manager'] + feature_data['JobRole_Manufacturing Director'] + \
                                feature_data['JobRole_Research Director'] + feature_data['JobRole_Sales Executive']
    # 开始工作年龄 = 年龄 - 总工作年限
    feature_data['开始工作年龄'] = feature_data['Age'] - feature_data['TotalWorkingYears']

    # 收入增长率 = 薪资涨幅百分比/在公司工作年限
    feature_data['入职收入增长率'] = feature_data['PercentSalaryHike'] / feature_data['YearsAtCompany']
    # 是否频繁换工作 = 曾工作过的公司数量 / 总工作年限
    feature_data['换工作频率'] = feature_data['NumCompaniesWorked'] / feature_data['TotalWorkingYears']
    # 在岗时间比例 = 在当前岗位工作年限 / 在公司工作年限
    feature_data['在岗时间比例'] = feature_data['YearsInCurrentRole'] / feature_data['YearsAtCompany']
    # 工作稳定指数 = 在公司工作年限 / 总工作年限
    feature_data['工作稳定指数'] = feature_data['YearsAtCompany'] / feature_data['TotalWorkingYears']
    # 搭档稳定指数 = 与当前经理共事年限 / 在当前岗位工作年限
    feature_data['搭档稳定指数'] = feature_data['YearsWithCurrManager'] / feature_data['YearsInCurrentRole']
    # 晋升停滞率 = 自上次晋升以来年限 / 在公司工作年限
    feature_data['晋升停滞率'] = feature_data['YearsSinceLastPromotion'] / feature_data['YearsAtCompany']

    # 工作年限 = 0 时，这些值为多少？ -------------与业务专家商量
    feature_data['入职收入增长率'] = feature_data['入职收入增长率'].replace([np.inf], 0)
    feature_data['换工作频率'] = feature_data['换工作频率'].replace([np.inf], 0)
    feature_data['在岗时间比例'] = feature_data['在岗时间比例'].fillna(0)
    feature_data['工作稳定指数'] = feature_data['工作稳定指数'].fillna(1)
    feature_data['晋升停滞率'] = feature_data['晋升停滞率'].fillna(0)
    feature_data['搭档稳定指数'] = feature_data['搭档稳定指数'].replace([np.inf], 0)
    feature_data['搭档稳定指数'] = feature_data['搭档稳定指数'].fillna(0)

    # 整体满意度指数 = （工作环境满意度 + 工作满意度 + 人际关系满意度 + 工作与生活平衡度） / 4
    feature_data['整体满意度指数'] = (feature_data['EnvironmentSatisfaction'] + feature_data['JobSatisfaction'] +
                                      feature_data['RelationshipSatisfaction'] +
                                      feature_data['WorkLifeBalance']) / 4
    # 满意度差值 = 工作满意度 - 工作环境满意度
    feature_data['满意度差值'] = feature_data['JobSatisfaction'] - feature_data['EnvironmentSatisfaction']

    # feature_columns = ['Age', 'MonthlyIncome', 'PerformanceRating', 'StockOptionLevel', 'DistanceFromHome', 'Education',
    #                    'TotalWorkingYears', 'YearsAtCompany', 'TrainingTimesLastYear', 'YearsInCurrentRole',
    #                    'YearsSinceLastPromotion', 'YearsWithCurrManager', 'BusinessTravel_Travel_Frequently',
    #                    'Gender_Female', 'MaritalStatus_Married', 'OverTime_Yes', 'isManager', '开始工作年龄',
    #                    '入职收入增长率', '换工作频率', '在岗时间比例', '工作稳定指数', '搭档稳定指数', '晋升停滞率',
    #                    '整体满意度指数', '满意度差值']
    feature_columns = ['StockOptionLevel', 'JobLevel', 'MaritalStatus_Single', 'OverTime_Yes', 'isManager', 'Age',
                       '换工作频率', 'TotalWorkingYears', 'YearsWithCurrManager', '入职收入增长率',
                       'YearsInCurrentRole', 'YearsAtCompany', '整体满意度指数', 'MonthlyIncome']
    feature_columns2 = ["MaritalStatus_Single", "OverTime_Yes", "isManager", "StockOptionLevel", "JobLevel",
                        "JobInvolvement", "JobSatisfaction", "EnvironmentSatisfaction", "WorkLifeBalance",
                        "TotalWorkingYears", "MonthlyIncome", "Age", "YearsAtCompany", "YearsInCurrentRole",
                        "YearsWithCurrManager", "整体满意度指数", "在岗时间比例", "YearsSinceLastPromotion",
                        "DistanceFromHome", "入职收入增长率", "换工作频率",'搭档稳定指数']
    gaba = ['Age', 'DistanceFromHome', 'Education',
            'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel',
            'JobSatisfaction', 'MonthlyIncome', 'NumCompaniesWorked',
            'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
            'StandardHours', 'StockOptionLevel', 'TotalWorkingYears',
            'TrainingTimesLastYear', 'WorkLifeBalance', 'YearsAtCompany',
            'YearsInCurrentRole', 'YearsSinceLastPromotion', 'YearsWithCurrManager',
            'BusinessTravel_Non-Travel', 'BusinessTravel_Travel_Frequently',
            'BusinessTravel_Travel_Rarely', 'Department_Human Resources',
            'Department_Research & Development', 'Department_Sales',
            'EducationField_Human Resources', 'EducationField_Life Sciences',
            'EducationField_Marketing', 'EducationField_Medical',
            'EducationField_Other', 'EducationField_Technical Degree',
            'Gender_Female', 'Gender_Male', 'JobRole_Healthcare Representative',
            'JobRole_Human Resources', 'JobRole_Laboratory Technician',
            'JobRole_Manager', 'JobRole_Manufacturing Director',
            'JobRole_Research Director', 'JobRole_Research Scientist',
            'JobRole_Sales Executive', 'JobRole_Sales Representative',
            'MaritalStatus_Divorced', 'MaritalStatus_Married',
            'MaritalStatus_Single', 'OverTime_Yes', '开始工作年龄',
            '入职收入增长率', '换工作频率', '在岗时间比例', '工作稳定指数', '搭档稳定指数', '晋升停滞率',
            '整体满意度指数', '满意度差值']

    data2 = [
        'Age', 'DistanceFromHome', 'Education',
        'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel',
        'JobSatisfaction', 'MonthlyIncome', 'NumCompaniesWorked',
        'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
        'StandardHours', 'StockOptionLevel', 'TotalWorkingYears',
        'TrainingTimesLastYear', 'WorkLifeBalance', 'YearsAtCompany',
        'YearsInCurrentRole', 'YearsSinceLastPromotion', 'YearsWithCurrManager',
        'BusinessTravel_Non-Travel',  # 保留一个BusinessTravel特征
        'Department_Human Resources',  # 保留一个Department特征
        'EducationField_Human Resources',  # 保留一个EducationField特征
        'Gender_Female',  # 保留一个Gender特征
        'JobRole_Healthcare Representative',  # 保留一个JobRole特征
        'MaritalStatus_Divorced',  # 保留一个MaritalStatus特征
        'OverTime_Yes',  # 保留一个OverTime特征
        '开始工作年龄',
        '入职收入增长率', '换工作频率', '在岗时间比例', '工作稳定指数', '搭档稳定指数', '晋升停滞率',
        '整体满意度指数', '满意度差值'
    ]
    data3 = [
        'Age', 'JobInvolvement', 'JobLevel',
        'JobSatisfaction', 'MonthlyIncome', 'StockOptionLevel',
        'TotalWorkingYears', 'YearsAtCompany', 'YearsInCurrentRole',
        'YearsWithCurrManager', 'JobRole_Sales Representative',
        'MaritalStatus_Single', 'OverTime_Yes', '开始工作年龄',
        '入职收入增长率', '换工作频率', '在岗时间比例', '工作稳定指数', '搭档稳定指数', '晋升停滞率',
        '整体满意度指数', '满意度差值'  # 仅保留一个OverTime特征
    ]

    return feature_data, gaba


def feature_extra(code=1):
    if code == 0:
        df = pd.read_csv('../../data/raw/test2.csv')
    else:
        df = pd.read_csv('../../data/raw/train.csv')
    cols = [
        "BusinessTravel", "Department", "Education", "EducationField",
        "EnvironmentSatisfaction", "Gender", "JobInvolvement", "JobLevel", "JobRole",
        "JobSatisfaction", "MaritalStatus", "OverTime", "PerformanceRating",
        "RelationshipSatisfaction", "TrainingTimesLastYear", "StockOptionLevel", "WorkLifeBalance"
    ]
    for col in cols:
        le = LabelEncoder()
        df[col] = le.fit_transform(df[col])
    feature_columns = [
        "BusinessTravel",
        "Department",
        "EducationField",
        "EnvironmentSatisfaction",
        "JobInvolvement",
        "JobLevel",
        "JobRole",
        "JobSatisfaction",
        "MaritalStatus",
        "OverTime",
        "StockOptionLevel",
        "WorkLifeBalance",
        "DistanceFromHome",
        "Age",
        "YearsInCurrentRole",
        "YearsSinceLastPromotion",
        "MonthlyIncome",
        "PercentSalaryHike",
        "YearsWithCurrManager",
        "NumCompaniesWorked",
        "TotalWorkingYears",
    ]
    x = df[feature_columns]
    y = df['Attrition']
    return x, y


def plot(x, y):
    df = pd.DataFrame({
        'isManager': x['isManager'],
        'Attrition': y
    })
    # 创建交叉表
    contingency_table = pd.crosstab(df['isManager'], df['Attrition'], margins=True)

    # 可视化
    # 分组计数
    grouped = df.groupby(['isManager', 'Attrition']).size().unstack()

    # 堆叠柱状图
    grouped.plot(kind='bar', stacked=True, color=['#5f9ea0', '#ff7373'])
    plt.title('Stacked Bar Chart of Attrition by isManager')
    plt.xlabel('Is Manager')
    plt.ylabel('Count')
    plt.legend(title='Attrition', labels=['No', 'Yes'])
    plt.xticks([0, 1], ['False', 'True'])
    plt.show()

    # grouped_pct = grouped.div(grouped.sum(1), axis=0)
    # grouped_pct.plot(kind='bar', stacked=True, color=['#5f9ea0', '#ff7373'])
    # plt.title('Percentage of Attrition by isManager')
    # plt.xlabel('Is Manager')
    # plt.ylabel('Percentage')
    # plt.legend(title='Attrition', labels=['No', 'Yes'])
    # plt.xticks([0, 1], ['False', 'True'])
    # plt.show()
    from scipy.stats import chi2_contingency

    chi2, p, dof, expected = chi2_contingency(contingency_table)
    print(f"Chi-square statistic: {chi2:.4f}")
    print(f"P-value: {p:.4f}")





class EmployeeAttritionFeatureEngineer(BaseEstimator, TransformerMixin):
    def __init__(self, handle_zeros=True, add_interactions=True):
        """
        人才流失预测特征工程优化器
        主要优化功能说明：
        特征缺陷修复：
        自动处理分母为零的情况
        修正收入增长率和换工作频率的计算公式
        新特征创建：
        添加更有效的特征如绝对收入增长（考虑实际收入基数）
        创建平均在职时长替代有问题的换工作频率
        计算管理岗位标志isManager
        交互特征：
        加班与满意度交互（加班_满意度差值）
        管理岗位与收入增长交互（管理岗_绝对增长）
        资历与满意度交互（资历_满意度）
        非线性变换：
        对晋升停滞年限进行对数变换（晋升停滞_log）
        收入增长率的平方根变换
        工作稳定指数的指数变换
        多重共线性处理：
        自动移除高度相关的原始特征
        提供推荐特征集get_feature_names()

        参数:
        handle_zeros: 是否处理分母为零的情况 (默认True)
        add_interactions: 是否添加交互特征 (默认True)
        """
        self.handle_zeros = handle_zeros
        self.add_interactions = add_interactions
        self.manager_roles = ['Manager', 'Director', 'Executive']
        self.manager_columns = ['JobRole_Manager',
                                'JobRole_Manufacturing Director',
                                'JobRole_Research Director',
                                'JobRole_Sales Executive']

    def fit(self, X, y=None):
        # 特征工程不需要在fit中学习参数
        return self

    def transform(self, X):
        X = X.copy()

        # 1. 修复原有特征缺陷
        X = self._fix_existing_features(X)

        # 2. 创建新特征
        X = self._create_new_features(X)

        # 3. 添加交互特征
        if self.add_interactions:
            X = self._add_interaction_features(X)

        # 4. 非线性变换
        X = self._apply_nonlinear_transforms(X)

        # 5. 清理临时特征
        X = self._clean_temporary_features(X)

        return X

    def _fix_existing_features(self, X):
        """修复有缺陷的特征计算"""
        # 处理分母为零的情况
        if self.handle_zeros:
            X['YearsAtCompany'] = X['YearsAtCompany'].replace(0, 1)
            X['YearsInCurrentRole'] = X['YearsInCurrentRole'].replace(0, 1)
            X['TotalWorkingYears'] = X['TotalWorkingYears'].replace(0, 1)

        # 修复收入增长率计算
        X['入职收入增长率'] = X['PercentSalaryHike'] / X['YearsAtCompany']

        # 修复换工作频率计算
        X['换工作频率'] = X['NumCompaniesWorked'] / (X['TotalWorkingYears'] + 1e-5)

        return X

    def _create_new_features(self, X):
        """创建新的有效特征"""
        # 1. 是否管理岗位
        X['isManager'] = X[self.manager_columns].any(axis=1).astype(int)

        # 2. 开始工作年龄
        X['开始工作年龄'] = X['Age'] - X['TotalWorkingYears']

        # 3. 绝对收入增长 (更有效的特征)
        X['绝对收入增长'] = X['MonthlyIncome'] * X['PercentSalaryHike'] / 10000

        # 4. 平均在职时长 (替代换工作频率)
        X['平均在职时长'] = X['TotalWorkingYears'] / (X['NumCompaniesWorked'] + 1)

        # 5. 在岗时间比例
        X['在岗时间比例'] = X['YearsInCurrentRole'] / X['YearsAtCompany']

        # 6. 工作稳定指数
        X['工作稳定指数'] = X['YearsAtCompany'] / (X['TotalWorkingYears'] + 1e-5)

        # 7. 搭档稳定指数
        X['搭档稳定指数'] = X['YearsWithCurrManager'] / X['YearsInCurrentRole']

        # 8. 晋升停滞率
        X['晋升停滞率'] = X['YearsSinceLastPromotion'] / X['YearsAtCompany']

        # 9. 整体满意度指数
        satisfaction_cols = ['EnvironmentSatisfaction',
                             'JobSatisfaction',
                             'RelationshipSatisfaction',
                             'WorkLifeBalance']
        X['整体满意度指数'] = X[satisfaction_cols].mean(axis=1)

        # 10. 满意度差值
        X['满意度差值'] = X['JobSatisfaction'] - X['EnvironmentSatisfaction']

        return X

    def _add_interaction_features(self, X):
        """添加重要的交互特征"""
        # 1. 加班与满意度交互
        X['加班_满意度差值'] = X['OverTime_Yes'] * X['满意度差值']

        # 2. 管理岗与收入增长交互
        X['管理岗_绝对增长'] = X['isManager'] * X['绝对收入增长']

        # 3. 工作年限与满意度交互
        X['资历_满意度'] = X['YearsAtCompany'] * X['整体满意度指数']

        # 4. 加班与工作生活平衡交互
        X['加班_工作生活平衡'] = X['OverTime_Yes'] * X['WorkLifeBalance']

        return X

    def _apply_nonlinear_transforms(self, X):
        """应用非线性变换"""
        # 1. 晋升停滞的对数变换
        X['晋升停滞_log'] = np.log1p(X['YearsSinceLastPromotion'])

        # 2. 收入增长率的平方根变换
        X['收入增长率_sqrt'] = np.sqrt(np.abs(X['入职收入增长率']))

        # 3. 工作稳定指数的指数变换
        X['工作稳定指数_exp'] = np.exp(X['工作稳定指数'])

        return X

    def _clean_temporary_features(self, X):
        """清理中间特征"""
        # 删除原始特征中可能引起多重共线性的特征
        redundant_cols = ['YearsInCurrentRole', 'YearsAtCompany', 'PercentSalaryHike']
        X.drop(columns=redundant_cols, errors='ignore', inplace=True)

        return X

    # def get_feature_names(self):
    #     """获取推荐特征集（可在特征选择后调用）"""
    #     return [
    #         'OverTime_Yes','MaritalStatus_Single', 'JobLevel','StockOptionLevel',
    #         'Age', '换工作频率', 'TotalWorkingYears', 'YearsWithCurrManager', '入职收入增长率',
    #         '绝对收入增长', '平均在职时长', 'YearsInCurrentRole', 'YearsAtCompany', '整体满意度指数',
    #         'MonthlyIncome', '加班_工作生活平衡', '收入增长率_sqrt', '资历_满意度'
    #     ]

# if __name__ == '__main__':
#     df = data_processing(code=1)
#     engineer = EmployeeAttritionFeatureEngineer()
#     df_engineered = engineer.transform(df)

# if __name__ == '__main__':
# x, y = data_processing()
# print(x.info())
# feature_data, feature_columns = feature_engineering(x)
# plot(feature_data, y)
# print(feature_data.info())
